# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/2 18:25 
@Author  : ZhangShenao 
@File    : 1.LangGraph基础组件上手.py 
@Desc    : LangGraph基础组件上手

在一个最基础LangGraph应用程序中,涵盖了3种基本组件:
1. 状态(State): 状态是图应用程序处理与交互的基础,是图中所有节点和边的输入
                它可以是一个dict(字典)或者Pydantic模型
                在LangGraph中,状态包括图的模式(数据结构)以及如何更新状态的归纳函数
                如果没有设置归纳函数,则每次节点的输出都会覆盖状态的原始数据
2. 节点(Node): 节点通常是Python函数(同步或异步)
                其中节点函数的第一个参数是state(状态)
                第二个参数是config(Runnable运行的配置)
                节点函数的返回值一般都是状态
                整个图的起点被称为开始节点,最后的终点被称为结束节点
3. 边(Edge): 边在图中定义了路由逻辑
                即不同节点之间是如何数据传递的、传递给谁、以及图节点从哪里开始、从哪里结束
                一个节点可以设置多条边
                如果有多条边,则下一条边连接的所有节点都会并行运行

LangGraph开发的基础步骤:
1. 初始化大语言模型和工具(ChatOpenAI、tools等等)
2. 定义图的状态State(TypedDict模型/Pydantic模型)
3. 用状态初始化图架构(StateGraph/MessageGraph)
4. 为图定义每一个节点(通过add_node方法为图添加节点)
5. 定义图的起点、终点和节点边(通过add_edge方法为图添加边)
6. 将图编译为Runnable可运行组件(通过graph.compile方法编译图)
7. 调用编译后的Runnable可运行组件执行图(通过graph.invoke方法调用图)
"""
from typing import TypedDict, Annotated

import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.runnables import RunnableConfig
from langgraph.constants import START, END
from langgraph.graph import add_messages, StateGraph

# 加载环境变量
dotenv.load_dotenv()


# 定义图状态
class GraphState(TypedDict):
    """
    图状态定义
    类型为Dict字典
    """

    messages: Annotated[list[str], add_messages]
    """
    状态字段——消息列表
    使用add_messages归纳函数,会将每个节点输出状态的messages字段聚合到列表中
    """

    node_name: str
    """状态字段——节点名称"""


# 创建LLM
llm = ChatZhipuAI(model="glm-4-plus")


# 定义Chat节点
# LangGraph的节点通常为一个处理函数
# 接收State状态作为输入,并输出处理后的状态
def chat_node(state: GraphState, config: RunnableConfig) -> GraphState:
    """聊天节点"""

    # 获取状态中存储的消息列表
    messages = state["messages"]

    # 调用LLM,获取输出
    resp = llm.invoke(messages)

    # 返回处理后的状态
    return {"messages": [resp], "node_name": "llm"}


# 创建Graph图架构
# 类型为StateGraph,并使用GraphState作为状态数据结构
graph_builder = StateGraph(GraphState)

# 添加节点
# 指定节点名称和节点函数
graph_builder.add_node(node="llm", action=chat_node)

# 添加边
# 为每个边指定起始节点和结束节点
graph_builder.add_edge(start_key=START, end_key="llm")
graph_builder.add_edge(start_key="llm", end_key=END)

# 将图编译为Runnable可运行组件
graph = graph_builder.compile()

# 定义图的初始状态
init_state = {"messages": ["你好，我是zsa，我喜欢写代码和打篮球"], "node_name": "start"}

# 调用图架构,获取结果
# 最终返回结果是一个GraphState状态
# 可以在LangSmith平台观测整个Graph结构的执行过程
result = graph.invoke(init_state)
print(result)
